import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt

try:
    # 加载 MNIST 数据集
    mnist = keras.datasets.mnist
    (train_images, train_labels), (test_images, test_labels) = mnist.load_data()

    # 数据预处理：将像素值缩放到 0 到 1 之间
    train_images = train_images / 255.0
    test_images = test_images / 255.0

    # 构建神经网络模型
    model = keras.Sequential([
        # 将 28x28 的二维图像展平为一维向量
        keras.layers.Flatten(input_shape=(28, 28)),
        # 全连接层，128 个神经元，使用 ReLU 激活函数
        keras.layers.Dense(128, activation='relu'),
        # 输出层，10 个神经元，对应 0 到 9 的 10 个数字类别，使用 softmax 激活函数
        keras.layers.Dense(10, activation='softmax')
    ])

    # 编译模型
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    # 训练模型
    model.fit(train_images, train_labels, epochs=5)

    # 评估模型
    test_loss, test_acc = model.evaluate(test_images, test_labels)
    print(f"测试集准确率: {test_acc}")

    # 进行预测
    predictions = model.predict(test_images)

    # 可视化预测结果
    def plot_image(i, predictions_array, true_label, img):
        predictions_array, true_label, img = predictions_array, true_label[i], img[i]
        plt.grid(False)
        plt.xticks([])
        plt.yticks([])

        plt.imshow(img, cmap=plt.cm.binary)

        predicted_label = np.argmax(predictions_array)
        if predicted_label == true_label:
            color = 'blue'
        else:
            color = 'red'

        plt.xlabel(f"{predicted_label} {100 * np.max(predictions_array):2.0f}% ({true_label})",
                   color=color)

    def plot_value_array(i, predictions_array, true_label):
        predictions_array, true_label = predictions_array, true_label[i]
        plt.grid(False)
        plt.xticks(range(10))
        plt.yticks([])
        thisplot = plt.bar(range(10), predictions_array, color="#777777")
        plt.ylim([0, 1])
        predicted_label = np.argmax(predictions_array)

        thisplot[predicted_label].set_color('red')
        thisplot[true_label].set_color('blue')



    # 显示前 15 张图像的预测结果
    num_rows = 5
    num_cols = 3
    num_images = num_rows * num_cols
    plt.figure(figsize=(2 * 2 * num_cols, 2 * num_rows))
    for i in range(num_images):
        plt.subplot(num_rows, 2 * num_cols, 2 * i + 1)
        plot_image(i, predictions[i], test_labels, test_images)
        plt.subplot(num_rows, 2 * num_cols, 2 * i + 2)
        plot_value_array(i, predictions[i], test_labels)
    plt.tight_layout()
    plt.show(block=True)  # 添加 block=True 参数

    input("按任意键继续...")
except Exception as e:
    print(f"出现异常: {e}")